# system modules ...
import os
import sys
import time
import datetime
import argparse
import json

# pytorch modules ..
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.utils.tensorboard as tensorboard    
torch.set_printoptions(precision=2)

# self modules ..
from utils.basetrainer import BaseTrainer


# <Class: RetinaFaceTrainer/>
class RetinaFaceTrainer(BaseTrainer):
    """
    Some Information about RetinaFaceTrainer:\n
    ...
    """
    # <Method: __init__/>
    def __init__(
        self, 
        train_loader=None, 
        val_loader=None, 
        eval_loader=None, 
        network=None, 
        criterion=None, 
        optimizer=None, 
        scheduler=None, 
        epochs=1000, 
        device='cpu', 
        tbx_writer=None, 
        checkpoints_folder='./checkpoints/', 
        priors=None,
        loc_weight=2.0
        ):
        super(RetinaFaceTrainer, self).__init__(train_loader, val_loader, eval_loader, network, criterion, optimizer, scheduler, epochs, device, tbx_writer, checkpoints_folder)
        self._priors = priors
        self._loc_weight = loc_weight
    # <Method: /__init__>

    # <Method: __get_val_batch_preds__/>
    def __get_val_batch_preds__(self, batch_val_data, *args, **kwargs):                
        locations, confidence, landmarks = self._network(batch_val_data[0].to(self._device))
        return locations, confidence, landmarks
    # <Method: /__get_val_batch_preds__>
    
    # <Method: __get_val_losses__/>
    def __get_val_losses__(self, batch_preds, batch_val_data, *args, **kwargs):
        loss_loc, loss_cls, _ = self._criterion(batch_preds, self._priors, batch_val_data[1])
        total_loss = loss_loc * self._loc_weight + loss_cls
        return { 'total_loss': total_loss, 'loss_loc': loss_loc, 'loss_cls': loss_cls }
    # <Method: /__get_val_losses__>              
    
    # <Method: __get_train_batch_preds__/>
    def __get_train_batch_preds__(self, batch_train_data, *args, **kwargs):
        locations, confidence, landmarks = self._network(batch_train_data[0].to(self._device))
        return locations, confidence, landmarks
    # <Method: /__get_train_batch_preds__>
    
    # <Method: __get_train_losses__/>
    def __get_train_losses__(self, batch_preds, batch_train_data, *args, **kwargs):
        loss_loc, loss_cls, loss_pts = self._criterion(batch_preds, self._priors, batch_train_data[1])
        total_loss = loss_loc * self._loc_weight + loss_cls + loss_pts
        return { 'total_loss': total_loss, 'loss_loc': loss_loc, 'loss_cls': loss_cls, 'loss_pts': loss_pts }
    # <Method: /__get_train_losses__>
    
    # <Method: __if_stop_trainning__/>
    def __if_stop_trainning__(self, *args, **kwargs):
        for param in self._optimizer.param_groups:
            lr=param["lr"]
        #   endfor
        if lr < kwargs['stop_lr']:
            # self.__eval__(-1, *args, **kwargs)
            return True
        #   endif
        return False
    # <Method: /__if_stop_trainning__>

    # <Method: __eval__/>
    def __eval__(self, epoch, *args, **kwargs):
        if epoch % 10 != 0 or epoch == 0:
            return None
        # end-if
        metrics = RetinaFaceTrainer.Evaluation(
            net = self._network, 
            box_priors = kwargs['box_priors'], 
            evalloader = self._eval_loader, 
            device = self._device, 
            confidence_threshold = kwargs['confidence_threshold'], 
            iou_threshold = kwargs['iou_threshold'], 
            if_show = kwargs['if_show']
            )
        return metrics     
    # <Method: /__eval__>

# <Class: /RetinaFaceTrainer>